#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "common.h"
#include "chunk.h"
#include "object.h"
#include "compiler.h"
#include "mem.h"
#include "scanner.h"

#ifdef DEBUG_PRINT_CODE
#include "debug.h"
#endif

typedef struct {
    Token current;
    Token previous;
    bool hadError;
    bool panicMode;
} Parser;

typedef enum {
    PREC_NONE,
    PREC_ASSIGNMENT,
    PREC_OR,
    PREC_AND,
    PREC_EQUALITY,
    PREC_COMPARISON,
    PREC_TERM,
    PREC_FACTOR,
    PREC_UNARY,
    PREC_CALL,
    PREC_PRIMARY,
} Precedence;

typedef void (*ParseFn)(bool canAssign);

typedef struct {
    ParseFn prefix;
    ParseFn infix;
    Precedence precedence;
} ParseRule;

static ParseRule rules[];

typedef struct {
    Token name;
    int depth;
    bool isCaptured;
} Local;

typedef struct {
    uint8_t index;
    bool isLocal;
} Upvalue;

typedef enum {
    TYPE_FUNCTION,
    TYPE_SCRIPT,
} FunctionType;

typedef struct RoxCompiler {
    struct RoxCompiler *parent;
    ObjFunction* function;
    FunctionType type;
    Local locals[UINT8_COUNT];
    int localCount;
    Upvalue upvalues[UINT8_COUNT];
    int scopeDepth;
} Compiler;

Parser parser;

Compiler *compiler = NULL;

static Chunk* currentChunk() {
    return &compiler->function->chunk;
}

static void errorAt(Token* token, const char* message) {
    if (parser.panicMode) {
        return;
    }
    parser.panicMode = true;
    fprintf(stderr, "[line %d] Error", token->line);

    if (token->type == TOKEN_EOF) {
        fprintf(stderr, " at end");
    } else if (token->type == TOKEN_ERROR) {
    } else {
        fprintf(stderr, " at '%.*s'", token->length, token->start);
    }

    fprintf(stderr, ": %s\n", message);
    parser.hadError = true;
}

static void errorMessage(const char* message) {
    errorAt(&parser.previous, message);
}

static void errorAtCurrent(const char* message) {
    errorAt(&parser.current, message);
}

static TokenType advance() {
    parser.previous = parser.current;

    while (1) {
        parser.current = scanToken();
        if (parser.current.type != TOKEN_ERROR) {
            break;
        }
        errorAtCurrent(parser.current.start);
    }
    return parser.previous.type;
}

static TokenType consume(TokenType type, const char* message) {
    if (parser.current.type == type) {
        return advance();
    }
    errorAtCurrent(message);
    return TOKEN_ERROR;
}

static bool match(TokenType type) {
    bool yes = false;
    if (parser.current.type == type) {
        advance();
        yes = true;
    }
    return yes;
}

static void emitByte(uint8_t byte) {
    writeChunk(currentChunk(), byte, parser.previous.line);
}

static void emitBytes(uint8_t byte1, uint8_t byte2) {
    emitByte(byte1);
    emitByte(byte2);
}

static void emitReturn() {
    emitByte(OP_NULL);
    emitByte(OP_RETURN);
}

static int emitJump(uint8_t instruction) {
    emitByte(instruction);
    emitByte(0xff);
    emitByte(0xff);
    return currentChunk()->size - 2;
}

static void emitLoop(int loopStart) {
    emitByte(OP_LOOP);

    int offset = currentChunk()->size - loopStart + 2;
    if (offset > UINT16_MAX) {
        errorMessage("Loop body too large!");
    }
    emitBytes((offset >> 8) & 0xff, offset & 0xff);
}

static uint8_t makeConstant(Value value) {
    int pos = addConstant(currentChunk(), value);
    if (pos > UINT8_MAX) {
        errorMessage("Too many constants in one chunk!");
    }
    return (uint8_t) pos;
}

static void emitConstant(Value value) {
    emitBytes(OP_CONSTANT, makeConstant(value));
}

static void patchJump(int offset) {
    int jump = currentChunk()->size - offset - 2;
    if (jump > UINT16_MAX) {
        errorMessage("Too much code to jump over");
    }

    uint8_t *code = currentChunk()->code + offset;
    *code = (jump >> 8) & 0xff;
    *(code + 1) = jump & 0xff;
}

static void initCompiler(Compiler *c, FunctionType type) {
    c->parent = compiler;
    c->function = NULL;
    c->type = type;
    c->localCount = 0;
    c->scopeDepth = 0;
    c->function = newFunction();
    compiler = c;

    if (type != TYPE_SCRIPT) {
        compiler->function->name = copyString(parser.previous.start, parser.previous.length);
    }

    Local* local = compiler->locals + compiler->localCount++;
    local->depth = 0;
    local->isCaptured = false;
    local->name.start = "";
    local->name.length = 0;
}

static ObjFunction* endCompiler() {
    emitReturn();
    ObjFunction* function = compiler->function;
#ifdef DEBUG_PRINT_CODE
    if (!parser.hadError) {
        disassembleChunk(&function->chunk, function->name != NULL
                                         ? function->name->str : "<script>");
    }
#endif
    compiler = compiler->parent;
    return function;
}

static void beginScope() {
    compiler->scopeDepth++;
}

static void endScope() {
    compiler->scopeDepth--;

    int n = compiler->localCount;
    while (n-- > 0 && (compiler->locals + n)->depth > compiler->scopeDepth) {
        if (compiler->locals[n].isCaptured) {
            emitByte(OP_CLOSE_UPVALUE);
        } else {
            emitByte(OP_POP);
        }
    }
    compiler->localCount = n + 1;
}

static ParseRule* getRule(TokenType type) {
    return &rules[type];
}

static void parseWithPrec(Precedence precedence) {
    TokenType type = advance();
    ParseFn prefixRule = getRule(type)->prefix;
    if (prefixRule == NULL) {
        errorMessage("Expect expression!");
        return;
    }

    bool canAssign = precedence <= PREC_ASSIGNMENT;
    prefixRule(canAssign);

    while (precedence <= getRule(parser.current.type)->precedence) {
        type = advance();
        ParseFn infixRule = getRule(type)->infix;
        infixRule(canAssign);
    }

    if (canAssign && match(TOKEN_EQUAL)) {
        errorMessage("Invalid assignment target");
    }
}

static void binary(bool canAssign) {
    TokenType operatorType = parser.previous.type;

    ParseRule *rule = getRule(operatorType);
    parseWithPrec(rule->precedence + 1);

    switch (operatorType) {
        case TOKEN_NOT_EQUAL: {
            emitBytes(OP_EQUAL, OP_NOT);
            break;
        }
        case TOKEN_EQUAL_EQUAL: {
            emitByte(OP_EQUAL);
            break;
        }
        case TOKEN_GREATER: {
            emitByte(OP_GREATER);
            break;
        }
        case TOKEN_GREATER_EQUAL: {
            emitBytes(OP_LESS, OP_NOT);
            break;
        }
        case TOKEN_LESS: {
            emitByte(OP_LESS);
            break;
        }
        case TOKEN_LESS_EQUAL: {
            emitBytes(OP_GREATER, OP_NOT);
            break;
        }
        case TOKEN_PLUS: {
            emitByte(OP_ADD);
            break;
        }
        case TOKEN_MINUS: {
            emitByte(OP_SUBTRACT);
            break;
        }
        case TOKEN_STAR: {
            emitByte(OP_MULTIPLY);
            break;
        }
        case TOKEN_SLASH: {
            emitByte(OP_DIVIDE);
            break;
        }
        default:
            return;
    }
}

static void literal(bool canAssign) {
    switch (parser.previous.type) {
        case TOKEN_FALSE: {
            emitByte(OP_FALSE);
            break;
        }
        case TOKEN_NULL: {
            emitByte(OP_NULL);
            break;
        }
        case TOKEN_TRUE: {
            emitByte(OP_TRUE);
            break;
        }
        default:
            return;
    }
}

static void expression();
static void statement();
static void declaration();
static void recovery();
static uint8_t parseVariable(const char* msg);
static void function(FunctionType type);

static void grouping(bool canAssign) {
    expression();
    consume(TOKEN_RIGHT_PAREN, "Expect ')' after expression!");
}

static void number(bool canAssign) {
    double value = strtod(parser.previous.start, NULL);
    emitConstant(NUMBER_VAL(value));
}

static void unary(bool canAssign) {
    TokenType operatorType = parser.previous.type;

    parseWithPrec(PREC_UNARY);

    switch (operatorType) {
        case TOKEN_NOT: {
            emitByte(OP_NOT);
            break;
        }
        case TOKEN_MINUS: {
            emitByte(OP_NEGATE);
            break;
        }
        default:
            return;
    }
}

static void string(bool canAssign) {
  ObjString *p = copyString(parser.previous.start + 1,
                       parser.previous.length - 2);

    emitConstant(OBJ_VAL(p));
}

static void expression() {
    parseWithPrec(PREC_ASSIGNMENT);
}

static void or_(bool canAssign) {
    int elseJump = emitJump(OP_JUMP_IF_FALSE);
    int endJump = emitJump(OP_JUMP);

    patchJump(elseJump);
    emitByte(OP_POP);

    parseWithPrec(PREC_OR);
    patchJump(endJump);
}

static void and_(bool canAssign) {
    int endJump = emitJump(OP_JUMP_IF_FALSE);

    emitJump(OP_POP);
    parseWithPrec(PREC_AND);

    patchJump(endJump);
}

static uint8_t argumentList() {
    uint8_t argCount = 0;
    if (parser.current.type != TOKEN_RIGHT_PAREN) {
        do {
            expression();
            if (argCount == 255) {
                errorMessage("Can't have more than 255 arguments");
            }
            argCount++;
        } while (match(TOKEN_COMMA));
    }
    consume(TOKEN_RIGHT_PAREN, "Expect ')' after arguments");
    return argCount;
}

static void call(bool canAssign) {
    uint8_t argCount = argumentList();
    emitBytes(OP_CALL, argCount);
}

static uint8_t identifierConstant(Token* token) {
    return makeConstant(OBJ_VAL(copyString(token->start, token->length)));
}

static void dot(bool canAssign) {
    consume(TOKEN_IDENTIFIER, "Expect property name after '.'!");
    uint8_t name = identifierConstant(&parser.previous);

    if (canAssign && match(TOKEN_EQUAL)) {
        expression();
        emitBytes(OP_SET_PROPERTY, name);
    } else {
        emitBytes(OP_GET_PROPERTY, name);
    }
}

static bool identifiersEqual(Token *l, Token *r) {
    if (l->length != r->length) return false;
    return memcmp(l->start, r->start, l->length) == 0;
}

static int resolveLocal(Compiler *c, Token *name) {
    int n = c->localCount;
    while (n-- > 0) {
        Local *local = c->locals + n;
        if (identifiersEqual(name, &local->name)) {
            if (local->depth == -1) {
                errorMessage("Can't read local variable in its own initializer!");
            }
            return n;
        }
    }
    return n;
}

static int addUpvalue(Compiler* c, uint8_t pos, bool isLocal) {
    int count = c->function->upvalueCount;

    for (int i = 0; i < count; i++) {
        Upvalue* upvalue = c->upvalues + i;
        if (upvalue->index == pos && upvalue->isLocal == isLocal) {
            return i;
        }
    }

    if (count == UINT8_COUNT) {
        errorMessage("Too many closure variables in function!");
        return 0;
    }
    Upvalue* upvalue = c->upvalues + count;
    upvalue->isLocal = isLocal;
    upvalue->index = pos;
    return c->function->upvalueCount++;
}

static int resolveUpvalue(Compiler* c, Token* name) {
    if (c->parent == NULL) {
        return -1;
    }

    int pos = resolveLocal(c->parent, name);
    if (pos != -1) {
        c->parent->locals[pos].isCaptured = true;
        return addUpvalue(c, (uint8_t) pos, true);
    }

    pos = resolveUpvalue(c->parent, name);
    if (pos != -1) {
        return addUpvalue(c, (uint8_t) pos, false);
    }
    return -1;
}

static void addLocal(Token* name) {
    if (compiler->localCount == UINT8_COUNT) {
        errorMessage("Too many local variables in function!");
        return;
    }
    Local *local = compiler->locals + compiler->localCount++;
    local->name = *name;
    local->depth = -1;
    local->isCaptured = false;
}

static void declareVariable() {
    if (compiler->scopeDepth == 0) {
        return;
    }
    Token *name = &parser.previous;
    int n = compiler->localCount;
    while (n-- > 0) {
        Local *local = compiler->locals + n;
        int depth = local->depth;
        if (depth != -1 && depth < compiler->scopeDepth) {
            break;
        }
        if (identifiersEqual(name, &local->name)) {
            errorMessage("Already variable with this name in this scope!");
        }
    }
    addLocal(name);
}

static uint8_t parseVariable(const char* msg) {
    consume(TOKEN_IDENTIFIER, msg);

    declareVariable();
    if (compiler->scopeDepth > 0) return 0;

    return identifierConstant(&parser.previous);
}

static void markInitialized() {
    if (compiler->scopeDepth == 0) {
        return;
    }
    compiler->locals[compiler->localCount - 1].depth = compiler->scopeDepth;
}

static void defineVariable(uint8_t index) {
    if (compiler->scopeDepth > 0) {
        markInitialized();
        return;
    }
    emitBytes(OP_DEFINE_GLOBAL, index);
}

static void namedVariable(Token token, bool canAssign) {
    uint8_t getOp, setOp;
    int arg = resolveLocal(compiler, &token);
    if (arg != -1) {
        getOp = OP_GET_LOCAL;
        setOp = OP_SET_LOCAL;
    } else if ((arg = resolveUpvalue(compiler, &token)) != -1) {
        getOp = OP_GET_UPVALUE;
        setOp = OP_SET_UPVALUE;
    } else {
        arg = identifierConstant(&token);
        getOp = OP_GET_GLOBAL;
        setOp = OP_SET_GLOBAL;
    }

    if (canAssign && match(TOKEN_EQUAL)) {
        expression();
        emitBytes(setOp, (uint8_t) arg);
    } else {
        emitBytes(getOp, (uint8_t) arg);
    }
}

static void variable(bool canAssign) {
    namedVariable(parser.previous, canAssign);
}

static void varDeclaration() {
    uint8_t global = parseVariable("Expect variable name!");
    if (match(TOKEN_EQUAL)) {
        expression();
    } else {
        emitByte(OP_NULL);
    }
    consume(TOKEN_SEMICOLON, "Expect ';' after variable declaration!");

    defineVariable(global);
}

static void funDeclaration() {
    uint8_t global = parseVariable("Expect function name!");
    markInitialized();
    function(TYPE_FUNCTION);
    defineVariable(global);
}

static void classDeclaration() {
    consume(TOKEN_IDENTIFIER, "Expect class name!");
    uint8_t nameConstant = identifierConstant(&parser.previous);
    declareVariable();

    emitBytes(OP_CLASS, nameConstant);
    defineVariable(nameConstant);

    consume(TOKEN_LEFT_BRACE, "Expect '{' before class body!");
    consume(TOKEN_RIGHT_BRACE, "Expect '}' after class body!");
}

static void declaration() {
    if (match(TOKEN_CLASS)) {
        classDeclaration();
    } else if (match(TOKEN_FUN)) {
        funDeclaration();
    } else if (match(TOKEN_VAR)) {
        varDeclaration();
    } else {
        statement();
    }

    if (parser.panicMode) {
        recovery();
    }
}

static void recovery() {
    parser.panicMode = false;
    while (parser.current.type != TOKEN_EOF) {
        if (parser.previous.type == TOKEN_SEMICOLON) {
            return;
        }
        switch (parser.current.type) {
            case TOKEN_CLASS:
            case TOKEN_FUN:
            case TOKEN_VAR:
            case TOKEN_FOR:
            case TOKEN_IF:
            case TOKEN_WHILE:
            case TOKEN_PRINT:
            case TOKEN_RETURN:
                return;
            default:
                break;
        }
        advance();
    }
}

static void printStatement() {
    expression();
    consume(TOKEN_SEMICOLON, "Expect ';' after value!");
    emitByte(OP_PRINT);
}

static void expressionStatement() {
    expression();
    consume(TOKEN_SEMICOLON, "Expect ';' after expression!");
    emitByte(OP_POP);
}

static void ifStatement() {
    consume(TOKEN_LEFT_PAREN, "Expect '(' after 'if'!");
    expression();
    consume(TOKEN_RIGHT_PAREN, "Expect ')' after if condition!");

    int thenPosition = emitJump(OP_JUMP_IF_FALSE);
    emitByte(OP_POP);
    statement();

    int elsePosition = emitJump(OP_JUMP);
    patchJump(thenPosition);
    emitByte(OP_POP);

    if (match(TOKEN_ELSE)) {
        statement();
    }
    patchJump(elsePosition);
}

static void returnStatement() {
    if (compiler->type == TYPE_SCRIPT) {
        errorMessage("Can't return from top-level code!");
    }
    if (match(TOKEN_SEMICOLON)) {
        emitReturn();
    } else {
        expression();
        consume(TOKEN_SEMICOLON, "Expect ';' after return value");
        emitByte(OP_RETURN);
    }
}

static void whileStatement() {
    int loopStart = currentChunk()->size;

    consume(TOKEN_LEFT_PAREN, "Expect '(' after 'while'!");
    expression();
    consume(TOKEN_RIGHT_PAREN, "Expect ')' after while condition!");

    int exitJump = emitJump(OP_JUMP_IF_FALSE);
    emitByte(OP_POP);

    statement();

    emitLoop(loopStart);

    patchJump(exitJump);
    emitByte(OP_POP);
}

static void forStatement() {
    beginScope();
    consume(TOKEN_LEFT_PAREN, "Expect '(' after 'for'!");
    if (match(TOKEN_SEMICOLON)) {
    } else if (match(TOKEN_VAR)) {
        varDeclaration();
    } else {
        expressionStatement();
    }

    int loopStart = currentChunk()->size;
    int exitJump = -1;
    if (!match(TOKEN_SEMICOLON)) {
        expression();
        consume(TOKEN_SEMICOLON, "Expect ';' after second expression of 'for'!");

        exitJump = emitJump(OP_JUMP_IF_FALSE);
        emitByte(OP_POP);
    }

    if (!match(TOKEN_RIGHT_PAREN)) {
        int bodyJump = emitJump(OP_JUMP);

        int incrementStart = currentChunk()->size;
        expression();
        emitByte(OP_POP);
        consume(TOKEN_RIGHT_PAREN, "Expect ')' after for clauses!");

        emitLoop(loopStart);
        loopStart = incrementStart;
        patchJump(bodyJump);
    }

    statement();

    emitLoop(loopStart);

    if (exitJump != -1) {
        patchJump(exitJump);
        emitByte(OP_POP);
    }
    endScope();
}

static void block() {
    TokenType type;
    while (type = parser.current.type, type != TOKEN_RIGHT_BRACE && type != TOKEN_EOF) {
        declaration();
    }
    consume(TOKEN_RIGHT_BRACE, "Expect '}' after block!");
}

static void function(FunctionType type) {
    Compiler c;
    initCompiler(&c, type);
    beginScope();

    consume(TOKEN_LEFT_PAREN, "Expect '(' after function name!");
    if (parser.current.type != TOKEN_RIGHT_PAREN) {
        do {
            compiler->function->arity++;
            if (compiler->function->arity > 255) {
                errorAtCurrent("Can't have more than 255 parameters!");
            }
            uint8_t paramConstant = parseVariable("Expect parameter name!");
            defineVariable(paramConstant);
        } while (match(TOKEN_COMMA));
    }
    consume(TOKEN_RIGHT_PAREN, "Expect ')' after function parameters!");

    consume(TOKEN_LEFT_BRACE, "Expect '{' before function body!");
    block();

    ObjFunction* func = endCompiler();
    uint8_t pos = makeConstant(OBJ_VAL(func));
    emitBytes(OP_CLOSURE, pos);

    for (int i = 0; i < func->upvalueCount; i++) {
        Upvalue* upvalue = c.upvalues + i;
        emitByte((uint8_t) (upvalue->isLocal ? 1 : 0));
        emitByte(upvalue->index);
    }
}

static void statement() {
    if (match(TOKEN_PRINT)) {
        printStatement();
    } else if (match(TOKEN_FOR)) {
        forStatement();
    } else if (match(TOKEN_IF)) {
        ifStatement();
    } else if (match(TOKEN_RETURN)) {
        returnStatement();
    } else if (match(TOKEN_WHILE)) {
        whileStatement();
    } else if (match(TOKEN_LEFT_BRACE)) {
        beginScope();
        block();
        endScope();
    } else {
        expressionStatement();
    }
}

static ParseRule rules[] = {
        [TOKEN_LEFT_PAREN]    = {grouping, call,   PREC_CALL},
        [TOKEN_RIGHT_PAREN]   = {NULL,     NULL,   PREC_NONE},
        [TOKEN_LEFT_BRACE]    = {NULL,     NULL,   PREC_NONE},
        [TOKEN_RIGHT_BRACE]   = {NULL,     NULL,   PREC_NONE},
        [TOKEN_COMMA]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_DOT]           = {NULL,     dot,    PREC_CALL},
        [TOKEN_MINUS]         = {unary,    binary, PREC_TERM},
        [TOKEN_PLUS]          = {NULL,     binary, PREC_TERM},
        [TOKEN_SEMICOLON]     = {NULL,     NULL,   PREC_NONE},
        [TOKEN_SLASH]         = {NULL,     binary, PREC_FACTOR},
        [TOKEN_STAR]          = {NULL,     binary, PREC_FACTOR},
        [TOKEN_NOT]           = {unary,    NULL,   PREC_NONE},
        [TOKEN_NOT_EQUAL]     = {NULL,     binary, PREC_NONE},
        [TOKEN_EQUAL]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_EQUAL_EQUAL]   = {NULL,     binary, PREC_EQUALITY},
        [TOKEN_GREATER]       = {NULL,     binary, PREC_COMPARISON},
        [TOKEN_GREATER_EQUAL] = {NULL,     binary, PREC_COMPARISON},
        [TOKEN_LESS]          = {NULL,     binary, PREC_COMPARISON},
        [TOKEN_LESS_EQUAL]    = {NULL,     binary, PREC_COMPARISON},
        [TOKEN_IDENTIFIER]    = {variable, NULL,   PREC_NONE},
        [TOKEN_STRING]        = {string,   NULL,   PREC_NONE},
        [TOKEN_NUMBER]        = {number,   NULL,   PREC_NONE},
        [TOKEN_AND]           = {NULL,     and_,   PREC_AND},
        [TOKEN_CLASS]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_ELSE]          = {NULL,     NULL,   PREC_NONE},
        [TOKEN_FALSE]         = {literal,  NULL,   PREC_NONE},
        [TOKEN_FOR]           = {NULL,     NULL,   PREC_NONE},
        [TOKEN_FUN]           = {NULL,     NULL,   PREC_NONE},
        [TOKEN_IF]            = {NULL,     NULL,   PREC_NONE},
        [TOKEN_NULL]          = {literal,  NULL,   PREC_NONE},
        [TOKEN_OR]            = {NULL,     or_,   PREC_OR},
        [TOKEN_PRINT]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_RETURN]        = {NULL,     NULL,   PREC_NONE},
        [TOKEN_SUPER]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_THIS]          = {NULL,     NULL,   PREC_NONE},
        [TOKEN_TRUE]          = {literal,  NULL,   PREC_NONE},
        [TOKEN_VAR]           = {NULL,     NULL,   PREC_NONE},
        [TOKEN_WHILE]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_ERROR]         = {NULL,     NULL,   PREC_NONE},
        [TOKEN_EOF]           = {NULL,     NULL,   PREC_NONE},
};

ObjFunction* compile(const char* source) {
    initScanner(source);
    Compiler jit;
    initCompiler(&jit, TYPE_SCRIPT);

    parser.hadError = false;
    parser.panicMode = false;

    advance();
    while (!match(TOKEN_EOF)) {
        declaration();
    }
    ObjFunction* function = endCompiler();
    return parser.hadError ? NULL : function;
}

void markCompilerRoots() {
    Compiler *c = compiler;
    while (c != NULL) {
        markObject((Obj *) compiler->function);
        c = c->parent;
    }
}
